from __future__ import division

import os
from xml.etree.ElementTree import Element, SubElement, ElementTree

import cv2
import numpy as np
import tensorflow as tf

from object_detection.utils import label_map_util
from utils.basic import cv_imread, cv_imwrite


def run_detection(sess, detection_graph, image_np):
    image_tensor = detection_graph.get_tensor_by_name('image_tensor:0')
    boxes = detection_graph.get_tensor_by_name('detection_boxes:0')
    scores = detection_graph.get_tensor_by_name('detection_scores:0')
    classes = detection_graph.get_tensor_by_name('detection_classes:0')
    num_detections = detection_graph.get_tensor_by_name('num_detections:0')
    (boxes, scores, classes, num_detections) = sess.run(
        [boxes, scores, classes, num_detections],
        feed_dict={
            image_tensor: image_np
        })
    boxes = np.squeeze(boxes)  # 从数组的形状中删除单维度条目，即把shape中为1的维度去掉
    classes = np.squeeze(classes).astype(np.int32)  # 从数组的形状中删除单维度条目，即把shape中为1的维度去掉 并将结果转换为整数型
    scores = np.squeeze(scores)
    return boxes, classes, scores


def load_model(checkpoint):
    """
    load the frozen graph of tensorflow as a detection model

    Parameters
    ----------
    checkpoint

    Returns
    -------

    """
    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
    detection_graph = tf.Graph()
    with detection_graph.as_default():
        od_graph_def = tf.GraphDef()
        with tf.gfile.GFile(checkpoint, 'rb') as fid:
            serialized_graph = fid.read()
            od_graph_def.ParseFromString(serialized_graph)
            tf.import_graph_def(od_graph_def, name='')
    return detection_graph


def get_file_name(file_name):
    """
    获取文件名，不包含扩展名

    Parameters
    ----------
    file_name: 文件完整路径

    Returns
    -------

    """
    (_, temp_filename) = os.path.split(file_name)
    (shot_name, _) = os.path.splitext(temp_filename)
    return shot_name


def is_overlap(rect1, rect2):
    """
    判断矩形是否存在重叠

    Parameters
    ----------
    rect1: 矩形1，形如[axis1_min, axis2_min, axis1_max, axis2_max]
    rect2: 矩形2，形如[axis1_min, axis2_min, axis1_max, axis2_max]

    Returns
    -------

    """
    return not ((rect1[0] >= rect2[2]) or
                (rect1[1] >= rect2[3]) or
                (rect1[2] <= rect2[0]) or
                (rect1[3] <= rect2[1]))


def get_overlap_area(rect1, rect2):
    """
    计算两个矩形的重叠部分的面积，如果不重叠返回0

    Parameters
    ----------
    rect1: 矩形1，形如[axis1_min, axis2_min, axis1_max, axis2_max]
    rect2: 矩形2，形如[axis1_min, axis2_min, axis1_max, axis2_max]

    Returns
    -------

    """
    if not is_overlap(rect1, rect2):
        return 0.0
    xmin = max(rect1[0], rect2[0])
    ymin = max(rect1[1], rect2[1])
    xmax = min(rect1[2], rect2[2])
    ymax = min(rect1[3], rect2[3])
    return (xmax - xmin) * (ymax - ymin)


def merge_boxes(all_boxes, all_classes, all_scores, overlap_percent):
    """
    合并重叠度过高的矩形框矩形框
    :param overlap_percent:
    :param all_boxes:
    :param all_classes:
    :param all_scores:
    :return: 无
    """
    need_del_idx = []  # 保存需要删除的box的索引
    for idx, box in enumerate(all_boxes):
        if idx not in need_del_idx:
            for id_2, box_2 in enumerate(all_boxes[idx + 1:], idx + 1):
                overlap_area = get_overlap_area(box, box_2)
                src_area = min((box[2] - box[0]) * (box[3] - box[1]), (box_2[2] - box_2[0]) * (box_2[3] - box_2[1]))
                overlap_rate = overlap_area / src_area
                if overlap_rate > overlap_percent:
                    if id_2 not in need_del_idx:
                        need_del_idx += [id_2]
    need_del_idx = sorted(need_del_idx, reverse=True)
    for i in need_del_idx:
        if i >= len(all_boxes):
            print("Error", i, len(all_boxes))
        del (all_boxes[i])
        del (all_classes[i])
        del (all_scores[i])
    return None


def write_result_2_xml(save_dir, image_name, image_shape, boxes, classes, scores, label_map_dict):
    """
    将检测结果写入xml文件
    :param scores: 所有box置信度
    :param save_dir: xml保存目录
    :param image_name: 图片文件名(xxx.xxx)
    :param image_shape: (height, width, depth)
    :param boxes:
    :param classes:
    :param label_map_dict: 编号到类别名称的映射字典
    """
    xml_path = os.path.join(save_dir, get_file_name(image_name) + ".xml")
    # with open(xml_path, 'wb') as f:
    root = Element("annotation")
    floder = SubElement(root, "floder")
    floder.text = "1"
    filename = SubElement(root, "filename")
    filename.text = image_name
    path = SubElement(root, "path")
    path.text = image_name
    source = SubElement(root, "source")
    database = SubElement(source, "database")
    database.text = "Unknown"
    size = SubElement(root, "size")
    width = SubElement(size, "width")
    width.text = str(image_shape[1])
    height = SubElement(size, "height")
    height.text = str(image_shape[0])
    depth = SubElement(size, "depth")
    depth.text = str(image_shape[2])
    segmented = SubElement(root, "segmented")
    segmented.text = "0"

    for _box, _class, _score in zip(boxes, classes, scores):
        object = SubElement(root, "object")
        name = SubElement(object, "name")
        name.text = str(label_map_dict[_class])
        pose = SubElement(object, "pose")
        pose.text = "Unspecified"
        truncated = SubElement(object, "truncated")
        truncated.text = "0"
        diffcult = SubElement(object, "diffcult")
        diffcult.text = "0"
        score = SubElement(object, "score")
        score.text = str(_score)
        region = SubElement(object, "region")
        region.text = "1"
        imageptr = SubElement(object, "imageptr")
        imageptr.text = "0"
        bndbox = SubElement(object, "bndbox")
        xmin = SubElement(bndbox, "xmin")
        xmin.text = str(int(_box[1] * image_shape[1]))
        ymin = SubElement(bndbox, "ymin")
        ymin.text = str(int(_box[0] * image_shape[0]))
        xmax = SubElement(bndbox, "xmax")
        xmax.text = str(int(_box[3] * image_shape[1]))
        ymax = SubElement(bndbox, "ymax")
        ymax.text = str(int(_box[2] * image_shape[0]))

    tree = ElementTree(root)
    tree.write(xml_path, encoding='utf-8')


class PredictObject:
    def __init__(self, model_path, image_dir, label_map_path, save_dir):
        self.model_path = model_path
        self.image_dir = image_dir
        self.label_map_path = label_map_path
        self.save_dir = save_dir

    def predict_single_image(self, image_path: str, sess, detection_graph, score, percent, label_map_dict_new):
        if image_path.endswith('.png') or image_path.endswith('.jpg') or image_path.endswith('.jpeg'):
            img = cv_imread(image_path)
            original_shape = img.shape
            original_image = img.copy()
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            img = np.expand_dims(img, axis=0)
            _boxes, _classes, _scores = run_detection(sess, detection_graph, img)

            boxes, classes, scores = [], [], []
            for box, cls, sc in zip(_boxes, _classes, _scores):
                if sc > score:
                    boxes.append(box)
                    classes.append(cls)
                    scores.append(sc)

            merge_boxes(boxes, classes, scores, percent)  # 对重合面积过大的box进行合并

            # 将检测结果写入xml
            image_base_name = os.path.basename(image_path)
            cv_imwrite(os.path.join(self.save_dir, image_base_name), original_image)
            write_result_2_xml(self.save_dir, image_base_name,
                               original_shape, boxes, classes, scores, label_map_dict_new)

    def predict_image(self, is_set, score=0.5, percent=0.7):
        """
        预测模型的精确度
        :param is_set:
        :param score: 置信度
        :param percent: iou重复度
        """
        # 加载模型
        detection_graph = load_model(self.model_path)
        label_map_dict = label_map_util.get_label_map_dict(self.label_map_path)
        label_map_dict_new = {value: key for key, value in label_map_dict.items()}
        # 配置文件
        config = tf.ConfigProto(device_count={"CPU": 4, "GPU": 1})
        config.gpu_options.allow_growth = True
        # 遍历文件夹下所有图片
        with detection_graph.as_default():
            with tf.Session(graph=detection_graph, config=config) as sess:
                if is_set:
                    for idx, image_name in enumerate(os.listdir(self.image_dir)):
                        self.predict_single_image(os.path.join(self.image_dir, image_name), sess, detection_graph,
                                                  score, percent, label_map_dict_new)
                else:
                    self.predict_single_image(self.image_dir, sess, detection_graph, score, percent, label_map_dict_new)


def predict_image(pb_path, label_map_path, test_image_dir, is_set, confidence_threshold, save_dir):
    """
    使用模型对图片进行预测
    :param pb_path: 模型路径
    :param label_map_path: label_map文件路径
    :param test_image_dir: 测试图片路径
    :param is_set: test_image_dir是单张图片路径还是文件夹路径，False为前者，True为后者
    :param confidence_threshold: 置信度阈值
    :param save_dir: 保存路径
    """
    if not os.path.exists(save_dir):
        os.mkdir(save_dir)
    pre = PredictObject(model_path=pb_path, image_dir=test_image_dir,
                        label_map_path=label_map_path,
                        save_dir=save_dir)
    pre.predict_image(is_set, score=confidence_threshold, percent=0.5)


if __name__ == "__main__":
    write_result_2_xml("", "货架", (100, 100, 3), [[0, 0, 1, 1], [1, 1, 1, 1]],
                       [1, 1], [0.9, 1.0], {1: "test"})
